import torch
import torch.nn as nn
import torch.nn.functional as F


class DepthwiseSeparableConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding=1):
        super(DepthwiseSeparableConv, self).__init__()
        self.depthwise = nn.Conv2d(
            in_channels,
            in_channels,
            kernel_size=kernel_size,
            padding=padding,
            groups=in_channels,
        )
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        x = F.relu(self.depthwise(x))
        x = self.pointwise(x)
        return x


class SqueezeExcitationLayer(nn.Module):
    def __init__(self, channel, reduction=16):
        super(SqueezeExcitationLayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid(),
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)


class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        self.down1 = DepthwiseSeparableConv(1, 64, 3)
        self.pool1 = nn.MaxPool2d(2)
        self.down2 = DepthwiseSeparableConv(64, 128, 3)
        self.pool2 = nn.MaxPool2d(2)
        self.se1 = SqueezeExcitationLayer(128)

        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.up_conv1 = DepthwiseSeparableConv(128, 64, 3)

        self.final_conv = nn.Conv2d(64, 1, 1)

    def forward(self, x):
        x1 = self.down1(x)
        x1p = self.pool1(x1)
        x2 = self.down2(x1p)
        x2 = self.se1(x2)
        x2 = self.pool2(x2)

        x = self.up1(x2)

        x = torch.cat([x, x1], dim=1)
        x = self.up_conv1(x)

        x = self.final_conv(x)
        return x


# 实例化模型
model = UNet()
print(model)
